from manimlib import *

__curr_dir = os.path.dirname(__file__)
if __curr_dir not in sys.path:
    sys.path.append(__curr_dir)
from MySpring import *

FONT_SIZE = 35
youngs_modules = 1  # KN/m2

REPEAT = 1


class DoubleSpring(Scene):
    def construct(self):
        """
        1. 生成一段[0,0,0]到[1,0,0]的弹簧
        2. 并压缩到[0.8,0,0]
        3. 显示F和u
        """
        axis = ThreeDAxes()
        frame = self.camera.frame
        frame.set_height(12)

        # ========================================================
        # ========================================================
        # ========================================================
        # 创建弹簧
        sprint_start = [0, 0, 0]
        sprint_end = [10, 0, 0]
        spring = Spring3D(sprint_start, sprint_end, axis)
        spring.center()  # 置于中间

        # ========================================================
        # ========================================================
        # ========================================================

        # 绘制文字
        level_0_text = Tex("Hook's Law: F=k\delta{L}")
        level_0_text.to_corner(UP)

        # ========================================================
        # ========================================================
        # ========================================================

        # 绘制宽度括号
        brace = always_redraw(Brace, spring, UP)
        brace.set_style(stroke_width=0.3)
        # 绘制括号标签
        brace_label2 = Tex("L = 0.00")
        # 更新位置
        brace_label2.always.next_to(brace, UP)

        # 更新数字
        def get_width():
            return spring.get_width() * 0.1

        spring_width_num = brace_label2.make_number_changable("0.00")
        spring_width_num.f_always.set_value(get_width)

        # ========================================================
        # ========================================================
        # ========================================================

        # 绘制箭头
        node1 = spring[1]
        node2 = spring[2]

        left_arrow = Arrow(start=node1.get_center() - np.array([1.5, 0, 0]), end=node1)
        left_arrow.set_style(stroke_color=RED)
        left_arrow.always.next_to(node1, LEFT)
        right_arrow = Arrow(start=node2.get_center() + np.array([1.5, 0, 0]), end=node2)
        right_arrow.set_style(stroke_color=RED)
        right_arrow.always.next_to(node2, RIGHT)

        # ========================================================
        # ========================================================
        # ========================================================

        # 绘制左F
        F_left_label = Tex("F_{1} = 0.00", t2c={"F_{1}": RED})

        # 更新位置
        F_left_label.always.next_to(left_arrow, LEFT)

        # 更新数字
        def get_value():
            return (0.1 * spring.get_width() - 1) * youngs_modules

        number = F_left_label.make_number_changable("0.00")
        number.f_always.set_value(get_value)

        # ========================================================
        # ========================================================
        # ========================================================

        # 绘制右F
        F_right_label = Tex("F_{2} = 0.00", t2c={"F_{2}": RED})
        # 更新位置
        F_right_label.always.next_to(right_arrow, RIGHT)

        # 更新数字
        def get_minus_value():
            return -1 * number.get_value()

        F_right_label.make_number_changable("0.00").f_always.set_value(get_minus_value)

        # ========================================================
        # ========================================================
        # ========================================================

        # 绘制左u
        u_left_label = Tex("\delta{u_{1}} = 0.00")
        # 更新位置
        u_left_label.always.next_to(F_left_label, DOWN)

        # 更新数字
        def get_half_value():
            return spring_width_num.get_value() / 2

        u_left_label.make_number_changable("0.00").f_always.set_value(get_half_value)

        # ========================================================
        # ========================================================
        # ========================================================

        # 绘制右u
        u_right_label = Tex("\delta{u_{2}} = 0.00")
        # 更新位置
        u_right_label.always.next_to(F_right_label, DOWN)

        # 更新数字
        def get_minus_half_value():
            return spring_width_num.get_value() / -2

        u_right_label.make_number_changable("0.00").f_always.set_value(get_minus_half_value)

        spring_group = Group(
            spring,
            brace,
            # brace_label1,
            brace_label2,
            # node 1
            left_arrow,
            F_left_label,
            u_left_label,
            # # node 2
            right_arrow,
            F_right_label,
            u_right_label,
        )

        # 在刚才的弹簧中加一个节点
        node3 = Sphere(radius=spring_group[0][1].radius, color=RED)
        spring_group.add(node3)
        # 更新位置
        node3.always.move_to(spring_group[0][0].get_center())

        mid_arrow31 = Arrow(
            start=node3,
            end=node3.get_center() - np.array([1.5, 0, 0]),
        )
        mid_arrow31.set_style(stroke_color=RED)
        always(mid_arrow31.next_to, node3, LEFT)

        # 绘制F31字
        F_31 = Tex(r"F_{31}", font_size=FONT_SIZE, color=RED)
        # 更新位置
        always(F_31.next_to, mid_arrow31, DOWN * 2)

        mid_arrow32 = Arrow(
            start=node3.get_center() + np.array([1.5, 0, 0]),
            end=node3,
        )
        mid_arrow32.set_style(stroke_color=RED)
        always(mid_arrow32.next_to, node3, RIGHT)

        # 绘制F32字
        F_32 = Tex(r"F_{32}", font_size=FONT_SIZE, color=RED)
        # 更新位置
        always(F_32.next_to, mid_arrow32, DOWN * 2)

        # ========================================================
        # ========================================================
        # ========================================================

        # 绘制u_3
        u_3 = Tex(r"\delta{u_{3}}", font_size=FONT_SIZE)
        # 更新位置
        always(u_3.next_to, node3, DOWN * 2)

        self.add(spring_group, level_0_text, node3, mid_arrow31, mid_arrow32, F_31, F_32)